import torch
import torchvision.models as models

#创建并训练一个模型
model=models.resent18(pretrained=True)
#___训练代码___

#保存整个模型
torch.save(model,'model_path')

#加载整个模型
loaded_model=torch.load('model.pth')

"""
仅保存模型参数（推荐方式）
"""
torch.save(model.state_dict(),'model__weights.pth')

#加载模型参数
model=models.resnet18() #必须先创建相同架构的模型
model.load_state_dict(torch.load('model_weights.path'))
model.eval() #设置为评估模式
#文件更小，更加灵活，兼容性更好

"""
保存和加载训练状态
"""

#保存检查点
checkpoint={
    'epoch':epoch,
    'model_state_dict':model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss':loss,
    #可以添加其他需要保存得信息
}

#加载保存点
checkpoint=torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch=checkpoint['epoch']
loss=checkpoint['loss']

model.eval()
#或者model.train()取决于你的需求

